{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python example for CMK-encryption in Azure AI Search\n",
    "\n",
    "This notebook provides sample script for [adding customer-managed key (CMK) encryption](https://learn.microsoft.com/azure/search/search-security-manage-encryption-keys) to objects on Azure AI Search."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "\n",
    "- [Azure AI Search](https://learn.microsoft.com/azure/search/search-create-service-portal)\n",
    "- [Azure Key Vault]()\n",
    "- [Azure Storage](https://learn.microsoft.com/azure/storage/common/storage-account-create) or [Azure Log Analytics](https://learn.microsoft.com/azure/azure-monitor/logs/quick-create-workspace?tabs=azure-portal) for data retention of audit logs.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sign in to Azure\n",
    "\n",
    "You might not need this step, but if downstream connections fail with a 401 during indexer pipeline execution, it could be because you're using the wrong tenant or subscription. You can avoid this issue by signing in from the command line, explicitly setting the tenant ID and choosing the right subscription.\n",
    "\n",
    "This section assumes you have the [Azure CLI](https://learn.microsoft.com/cli/azure/authenticate-azure-cli-interactively).\n",
    "\n",
    "1. Open a command line prompt.\n",
    "\n",
    "1. Run this command to get a list of Azure tenants: `az account tenant list`\n",
    "\n",
    "1. If you have multiple tenants, set the tenant: `az login --tenant <YOUR-TENANT_ID>`\n",
    "\n",
    "If you have multiple subscriptions, a list is provided so that you can select one."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a virtual environment in Visual Studio Code\n",
    "\n",
    "Create a virtual environment so that you can install the dependencies in isolation.\n",
    "\n",
    "1. In Visual Studio Code, open the folder containing tutorial-rag.ipynb.\n",
    "\n",
    "1. Press Ctrl-shift-P to open the command palette, search for \"Python: Create Environment\", and then select `Venv` to create a virtual environment in the current workspace.\n",
    "\n",
    "1. Select Tutorial-RAG\\tutorial-rag-requirements.txt for the dependencies.\n",
    "\n",
    "It takes several minutes to create the environment. When the environment is ready, continue to the next step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install python-dotenv\n",
    "! pip install azure-core\n",
    "! pip install azure-search-documents\n",
    "! pip install azure-storage-blob\n",
    "! pip install azure-identity\n",
    "! pip install openai\n",
    "! pip install aiohttp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CREATE INDEX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set variables\n",
    "AZURE_SEARCH_SERVICE: str = \"<PUT_YOUR_AZURE_SEARCH_SERVICE_URL_HERE>\"\n",
    "AZURE_KEY_VAULT_NAME: str = \"<PUT_YOUR_AZURE_KEY_VAULT_NAME_HERE>\"\n",
    "AZURE_KEY_VAULT_URI: str = \"<PUT_YOUR_AZURE_KEY_VAULT_URI_HERE>\"\n",
    "AZURE_KEY_VAULT_VERSION: str = \"<PUT-YOUR-AZURE_KEY_VAULT_VERSION_HERE>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " test-cmk-index-qs created\n"
     ]
    }
   ],
   "source": [
    "from azure.search.documents.indexes import SearchIndexClient\n",
    "from azure.search.documents.indexes.models import (\n",
    "    SimpleField,\n",
    "    SearchFieldDataType,\n",
    "    SearchableField,\n",
    "    SearchIndex,\n",
    "    SearchResourceEncryptionKey\n",
    ")\n",
    "from azure.identity import DefaultAzureCredential\n",
    "\n",
    "credential = DefaultAzureCredential()\n",
    "\n",
    "# Create a search index  \n",
    "index_name = \"test-cmk-index\"\n",
    "index_client = SearchIndexClient(endpoint=AZURE_SEARCH_SERVICE, credential=credential)  \n",
    "fields = [\n",
    "        SimpleField(name=\"Id\", type=SearchFieldDataType.String, key=True),\n",
    "        SearchableField(name=\"Description\", type=SearchFieldDataType.String)\n",
    "    ]\n",
    "\n",
    "scoring_profiles = []\n",
    "suggester = []\n",
    "encryption_key = SearchResourceEncryptionKey(\n",
    "    key_name=AZURE_KEY_VAULT_NAME,\n",
    "    key_version=AZURE_KEY_VAULT_VERSION,\n",
    "    vault_uri=AZURE_KEY_VAULT_URI\n",
    ")\n",
    "\n",
    "# Create the search index=\n",
    "index = SearchIndex(name=index_name, fields=fields, encryption_key=encryption_key)\n",
    "result = index_client.create_or_update_index(index)\n",
    "print(f' {result.name} created')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GET INDEX DEFINITION\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index_name = \"test-cmk-index\"\n",
    "index_client = SearchIndexClient(endpoint=AZURE_SEARCH_SERVICE, credential=credential)  \n",
    "\n",
    "result = index_client.get_index(index_name)  \n",
    "print(f\"{result}\")  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LOAD INDEX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Upload of new document succeeded: True\n"
     ]
    }
   ],
   "source": [
    "from azure.search.documents import SearchClient\n",
    "\n",
    "# Create a documents payload\n",
    "documents = [\n",
    "    {\n",
    "    \"@search.action\": \"upload\",\n",
    "    \"Id\": \"1\",\n",
    "    \"Description\": \"The hotel is ideally located on the main commercial artery of the city in the heart of New York. A few minutes away is Time's Square and the historic centre of the city, as well as other places of interest that make New York one of America's most attractive and cosmopolitan cities.\"\n",
    "    },\n",
    "    {\n",
    "    \"@search.action\": \"upload\",\n",
    "    \"Id\": \"2\",\n",
    "    \"Description\": \"The hotel is situated in a  nineteenth century plaza, which has been expanded and renovated to the highest architectural standards to create a modern, functional and first-class hotel in which art and unique historical elements coexist with the most modern comforts.\"\n",
    "    },\n",
    "    {\n",
    "    \"@search.action\": \"upload\",\n",
    "    \"Id\": \"3\",\n",
    "    \"Description\": \"The hotel stands out for its gastronomic excellence under the management of William Dough, who advises on and oversees all of the Hotel's restaurant services.\"\n",
    "    },\n",
    "    {\n",
    "    \"@search.action\": \"upload\",\n",
    "    \"Id\": \"4\",\n",
    "    \"Description\": \"The hotel is located in the heart of the historic center of Sublime in an extremely vibrant and lively area within short walking distance to the sites and landmarks of the city and is surrounded by the extraordinary beauty of churches, buildings, shops and monuments. Sublime Cliff is part of a lovingly restored 1800 palace.\"\n",
    "    }\n",
    "]\n",
    "\n",
    "search_client = SearchClient(endpoint=AZURE_SEARCH_SERVICE, index_name=index_name, credential=credential)\n",
    "try:\n",
    "    result = search_client.upload_documents(documents=documents)\n",
    "    print(\"Upload of new document succeeded: {}\".format(result[0].succeeded))\n",
    "except Exception as ex:\n",
    "    print (ex.message)\n",
    "\n",
    "    index_client = SearchClient(endpoint=AZURE_SEARCH_SERVICE, credential=credential)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QUERY ENCRYPTED CONTENT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Score: 0.6130029\n",
      "Id: 4\n",
      "Description: The hotel is located in the heart of the historic center of Sublime in an extremely vibrant and lively area within short walking distance to the sites and landmarks of the city and is surrounded by the extraordinary beauty of churches, buildings, shops and monuments. Sublime Cliff is part of a lovingly restored 1800 palace.\n",
      "Score: 0.26286605\n",
      "Id: 1\n",
      "Description: The hotel is ideally located on the main commercial artery of the city in the heart of New York. A few minutes away is Time's Square and the historic centre of the city, as well as other places of interest that make New York one of America's most attractive and cosmopolitan cities.\n"
     ]
    }
   ],
   "source": [
    "from azure.search.documents import SearchClient\n",
    "\n",
    "query = \"historic\"  \n",
    "\n",
    "search_client = SearchClient(endpoint=AZURE_SEARCH_SERVICE, credential=credential, index_name=index_name)\n",
    "  \n",
    "results = search_client.search(  \n",
    "    query_type='simple',\n",
    "    search_text=query, \n",
    "    select=[\"Id\", \"Description\"],\n",
    "    include_total_count=True\n",
    "    )\n",
    "  \n",
    "for result in results:  \n",
    "    print(f\"Score: {result['@search.score']}\")\n",
    "    print(f\"Id: {result['Id']}\")\n",
    "    print(f\"Description: {result['Description']}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
